package net.xuele.learn.flink.book;

import net.xuele.learn.flink.book.utils.SensorReading;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.util.Collector;

import java.util.concurrent.atomic.LongAdder;

/**
 * @Author patrick
 * @Date 2023/7/13 15:06
 * @Description
 */
public class HighTempCounter implements CheckpointedFunction, FlatMapFunction<SensorReading, Tuple3<String, Long, Long>> {

    private Double threshold;

    public HighTempCounter(Double threshold) {
        this.threshold = threshold;
    }

    // 在本地用于存储算子实例高温数目的变量
    long opHighTempCnt = 0;
    ValueState<Long> keyedCntState;
    ListState<Long> opCntState;


    @Override
    public void flatMap(SensorReading value, Collector<Tuple3<String, Long, Long>> out) throws Exception {
        if (value.temperature > threshold) {
            opHighTempCnt += 1;
            // 注意keyedCntState空指针
            long keyHighTempCnt = keyedCntState.value() + 1;
            keyedCntState.update(keyHighTempCnt);
            out.collect(Tuple3.of(value.id, keyHighTempCnt, opHighTempCnt));
        }
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        // 初始化话键值分区状态
        ValueStateDescriptor<Long> keyedCnt = new ValueStateDescriptor<>("keyedCnt", Long.class);
        keyedCntState = context.getKeyedStateStore().getState(keyedCnt);

        // 初始化算子状态
        ListStateDescriptor<Long> opCnt = new ListStateDescriptor<>("opCnt", Long.class);
        opCntState = context.getOperatorStateStore().getListState(opCnt);
        // 利用算子状态初始化本地遍历
        for (Long cnt : opCntState.get()) {
            // 可以看到ListState每个元素存储的都是一个”分片“值，类似于LongAdder里面的Cell（LongAdder的sum方法是Cell的综合）
            // 这里可以理解为每个算子实例分得其中的一份
            opHighTempCnt += cnt;
        }
    }


    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        // 对于每个算子实例，将当前的实例的状态进行添加
        // 这就是为什么不直接使用ValueState，而是使用ListState；若使用ValueState，则无法应对算子实例并行度的变化
        opCntState.clear();
        opCntState.add(opHighTempCnt);
    }
}
